package com.zhouchj.elasticsearch.update;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ReflectUtil;
import com.zhouchj.elasticsearch.basic.SFunction;
import com.zhouchj.elasticsearch.constant.EsConstant;
import com.zhouchj.elasticsearch.index.EsIndexOperate;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.springframework.data.elasticsearch.core.ElasticsearchRestTemplate;
import org.springframework.data.elasticsearch.core.RefreshPolicy;
import org.springframework.data.elasticsearch.core.ScriptType;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.*;

import java.lang.reflect.Field;
import java.util.*;
import java.util.stream.Collectors;

/**
 * ES链式更新
 *
 * @author zhouchj
 * @date 2023/8/18
 */
@Slf4j
public class EsLambdaUpdateChainWrapper<T> extends EsUpdate<T, SFunction<T, ?>, EsLambdaUpdateChainWrapper<T>> {

    private static final Integer BATCH_SIZE = EsConstant.BATCH_SIZE;

    private final EsIndexOperate<T> esIndexOperate;

    public EsLambdaUpdateChainWrapper(IndexCoordinates index,
                                      Class<T> clazz,
                                      ElasticsearchRestTemplate elasticsearchRestTemplate,
                                      RestHighLevelClient client) {
        super(clazz, index, elasticsearchRestTemplate);
        esIndexOperate = new EsIndexOperate<>(index, clazz, elasticsearchRestTemplate,client);
    }


    /**
     * 更新(终止操作)
     * 优先根据id更新，如果id没有设置则根据条件更新
     */
    public void update() {
        if (CollUtil.isEmpty(param)) {
            return;
        }
        if (CollUtil.isNotEmpty(ids)) {
            List<UpdateParam> updateParamList = ids.stream()
                    .map(id -> new UpdateParam(id, param))
                    .collect(Collectors.toList());
            update(updateParamList);
            return;
        }
        updateByQuery();
    }

    /**
     * 根据传入的实体类更新(终止操作)
     * null字段不更新 空字符串字段会更新
     *
     * @param entity es实体类
     */
    public void updateById(T entity) {
        this.updateBatchById(Collections.singletonList(entity));
    }

    /**
     * 根据传入的实体集合更新(终止操作)
     * null字段不更新 空字符串字段会更新
     *
     * @param list es实体集合
     */
    public void updateBatchById(List<T> list) {
        Field[] fields = ReflectUtil.getFields(clazz);
        List<UpdateParam> updateParamList = new ArrayList<>();
        for (T entity : list) {
            UpdateParam updateParam = reflectEntity(entity, fields);
            updateParamList.add(updateParam);
        }
        this.update(updateParamList);
        clear();
    }


    /**
     * 根据传入的实体类保存(终止操作)
     *
     * @param entity es实体类
     */
    public void save(T entity) {
        if (Objects.isNull(entity)) {
            return;
        }
        this.saveBatch(Collections.singletonList(entity));
    }

    /**
     * 根据传入的实体集合保存(终止操作)
     *
     * @param list es实体集合
     */
    public void saveBatch(List<T> list) {
        esIndexOperate.createIfExists();
        if (CollUtil.isEmpty(list)) {
            return;
        }
        // 单次保存
        if (list.size() <= BATCH_SIZE) {
            this.saveList(list);
            return;
        }
        // 分批次保存
        List<List<T>> split = CollUtil.split(list, BATCH_SIZE);
        for (List<T> splitList : split) {
            this.saveList(splitList);
        }
        clear();
    }

    /**
     * 删除(终止操作)
     * 优先根据id更新，如果id没有设置则根据条件更新
     */
    public void delete() {
        if (CollUtil.isNotEmpty(ids)) {
            deleteBatch(ids);
            return;
        }
        // 根据查询条件
        elasticsearchRestTemplate.delete(getQuery(), clazz, index);
        clear();
    }

    /**
     * 根据id集合批量删除(终止操作)
     *
     * @param idList id集合
     */
    public void deleteBatch(List<Long> idList) {
        Criteria criteria = new Criteria("id").in(idList);
        Query query = new CriteriaQuery(criteria);
        // 执行删除
        elasticsearchRestTemplate.delete(query, clazz, index);
        clear();
    }


    /**
     * 清除数据，让该类可以重复利用
     */
    public void clear() {
        boolQueryBuilder = new BoolQueryBuilder();
        builder = new NativeSearchQueryBuilder();
        ids.clear();
        param.clear();
    }


    /**
     * 保存一批数据
     *
     * @param list 具体数据
     */
    private void saveList(List<T> list) {
        // 构建IndexQuery对象列表
        List<IndexQuery> indexQueries = list.stream()
                .map(data -> new IndexQueryBuilder()
                        .withObject(data)
                        .build())
                .collect(Collectors.toList());
        BulkOptions bulkOptions = BulkOptions.builder()
                .withRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
                .build();
        elasticsearchRestTemplate.bulkIndex(indexQueries, bulkOptions, index);
    }


    /**
     * 反射实体类获取更新参数
     *
     * @param entity 实体类
     * @param fields 实例类字段
     * @return UpdateParam
     */
    private UpdateParam reflectEntity(T entity, Field[] fields) {
        Object id = null;
        Map<String, Object> map = new HashMap<>(fields.length);
        for (Field field : fields) {
            if (isId(field)) {
                id = ReflectUtil.getFieldValue(entity, field);
                continue;
            }
            Object fieldValue = ReflectUtil.getFieldValue(entity, field);
            if (Objects.nonNull(fieldValue)) {
                String fieldName = this.convertToEsFieldName(field);
                map.put(fieldName, fieldValue);
            }
        }
        if (Objects.isNull(id)) {
            throw new RuntimeException("id不存在或者为null");
        }
        Long longId = Long.valueOf(String.valueOf(id));
        return new UpdateParam(longId, map);
    }


    /**
     * 更新
     *
     * @param updateParamList 更新参数
     */
    private void update(List<UpdateParam> updateParamList) {
        // 拼接入参
        List<UpdateQuery> updateQueries = updateParamList
                .stream()
                .map(r -> UpdateQuery.builder(String.valueOf(r.getId()))
                        .withScriptType(ScriptType.INLINE)
                        .withScript(getUpdateScript(r.getParam()))
                        .withParams(r.getParam())
                        .build())
                .collect(Collectors.toList());
        // 拆分数组 多次更新
        List<List<UpdateQuery>> split = CollUtil.split(updateQueries, BATCH_SIZE);
        for (List<UpdateQuery> queries : split) {
            BulkOptions options = BulkOptions.builder()
                    .withRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
                    .build();
            elasticsearchRestTemplate.bulkUpdate(queries, options, index);
        }
        clear();
    }


    /**
     * 根据查询更新
     */
    private void updateByQuery() {
        UpdateQuery query = UpdateQuery.builder(getQuery())
                .withBatchSize(BATCH_SIZE)
                .withScriptType(ScriptType.INLINE)
                .withScript(getUpdateScript(param))
                .withParams(param)
                .withRefreshPolicy(RefreshPolicy.IMMEDIATE)
                .build();
        elasticsearchRestTemplate.updateByQuery(query, index);
        clear();
    }

    /**
     * 获取更新的脚本
     * 根据条件更新必须使用script
     *
     * @param param 需要更新的内容
     * @return 脚本
     */
    private String getUpdateScript(Map<String, Object> param) {
        StringBuilder sb = new StringBuilder();
        param.forEach((k, v) -> sb.append("ctx._source.")
                .append(k)
                .append("=params.")
                .append(k)
                .append(";"));
        return sb.toString();
    }

    @Data
    @AllArgsConstructor
    private static class UpdateParam {
        private Long id;

        private Map<String, Object> param;
    }

}
